import numpy as np
import matplotlib.pyplot as plt

m_x = np.loadtxt('../data/mnist_x', delimiter=' ')
m_y = np.loadtxt('../data/mnist_y')

# 划分训练集和测试集
ratio = 0.8
split = int(len(m_x) * ratio)
x_train,x_test = m_x[:split],m_x[split:]
y_train,y_test = m_y[:split],m_y[split:]

# 计算距离
def distance(a,b):
    return np.sqrt(np.sum((a-b)**2))

class KNN:
    def __init__(self,k=3):
        self.k = k
    def fit(self,x_train,y_train):
        self.x_train = x_train
        self.y_train = y_train
    def get_knn_indices(self,x):
        dis = list(map(lambda a: distance(a, x), self.x_train))
        knn_indices = np.argsort(dis)
        knn_indices = knn_indices[:self.k]
        return knn_indices
    def get_label(self,x):
        knn_indices = self.get_knn_indices(x)
        label_statistic = np.zeros(shape=10)
        for index in knn_indices:
            label = int(y_train[index])
            label_statistic[label] += 1
        return np.argmax(label_statistic)
    def predict(self,x_test):
        prediction = np.zeros(shape=len(x_test))
        for indexx,data in enumerate(x_test):
            prediction[indexx] = self.get_label(data)
        return prediction

for k in range(1,21):
    knn = KNN(k)
    knn.fit(x_train,y_train)
    accuracy = np.mean(y_test == knn.predict(x_test))
    print(k,":",accuracy)
